import numpy as np
from numpy.lib.stride_tricks import as_strided as ast
from sklearn.decomposition import PCA


def DCT2(X):
    A = np.float32(np.zeros_like(X))
    w = len(X)
    h = len(X.transpose())
    for i in range(w):
        for j in range(h):
            if i == 0:
                a = np.sqrt(1. / w)
            else:
                a = np.sqrt(2. / w)
            A[i, j] = a * np.cos(np.pi * (j + 0.5) * i / w)
            if abs(A[i, j]) > 0.4:
                A[i, j] = 0
    Y = np.dot(np.dot(A, X), A.transpose())
    return Y


def IDCT2(X):
    A = np.float32(np.zeros_like(X))
    w = len(X)
    h = len(X.transpose())
    for i in range(w):
        for j in range(h):
            if i == 0:
                a = np.sqrt(1. / w)
            else:
                a = np.sqrt(2. / w)
            A[i, j] = a * np.cos(np.pi * (j + 0.5) * i / w)
            if abs(A[i, j]) > 0.4:
                A[i, j] = 0
    Y = np.dot(np.dot(A, X), A.transpose())

    X1 = np.dot(np.dot(A.transpose(), Y), A)
    return X1


def block_view(A, block=(3, 3)):
    """Provide a 2D block view to 2D array. No error checking made.
    Therefore meaningful (as implemented) only for blocks strictly
    compatible with the shape of A."""
    # simple shape and strides computations may seem at first strange
    # unless one is able to recognize the 'tuple additions' involved ;-)
    shape = (A.shape[0] / block[0], A.shape[1] / block[1]) + block
    strides = (block[0] * A.strides[0], block[1] * A.strides[1]) + A.strides
    return ast(A, shape=shape, strides=strides), shape


def blkproc(A, block=(8, 8), func=None):
    view, shape = block_view(A=A, block=block)
    viewed = np.zeros_like(view)
    whole = None
    v = None
    for i in range(shape[0]):
        if i == 1:
            whole = v
        elif i >= 1:
            whole = np.vstack((whole, v))
        v = None
        for j in range(shape[1]):
            viewed[i, j] = func(view[i, j])
            if j == 0:
                v = viewed[i, j]
            else:
                v = np.hstack((v, viewed[i, j]))
    whole = np.vstack((whole, v))
    return whole


def haar_matrix(size):
    level = int(np.ceil(np.log2(size)))
    H = np.array([1.])[:, None]
    NC = 1. / np.sqrt(2.)
    LP = np.array([1., 1.])[:, None]
    HP = np.array([1., -1.])[:, None]
    for i in range(level):
        H = NC * np.hstack((np.kron(H, LP), np.kron(np.eye(len(H)), HP)))
    H = H.T
    return H


def dwt2d(x):
    H = haar_matrix(x.shape[0])
    return np.dot(np.dot(H, x), H.T)


def idwt2d(x):
    H = haar_matrix(x.shape[0])
    return np.dot(np.dot(H.T, x), H)


def wthresh(a, thresh):
    # Soft threshold
    res = np.abs(a) - thresh
    return np.sign(a) * ((res > 0) * res)


def wavelet_haar(A, thr):
    wt = dwt2d(A)
    wt = wthresh(wt, thr)
    return wt


def inv_wavelet_haar(A, thr):
    wt = dwt2d(A)
    wt = wthresh(wt, thr)
    rs = idwt2d(wt)
    return rs


def PCAdenoising(A, block, keep_dim):
    A = np.asarray(A, np.float32)
    view, shape = block_view(A=A, block=block)
    A_flat = []
    for i in range(shape[0]):
        for j in range(shape[1]):
            A_flat.append(np.reshape(view[i, j], (-1)))
    A_flat = np.asarray(A_flat, dtype=np.float32)
    (num_samples, fea_dim) = A_flat.shape
    fea_means = []
    for d in range(fea_dim):
        fea_mean = np.mean(A_flat[:, d])
        fea_means.append(fea_mean)
        A_flat[:, d] -= fea_mean
    pca = PCA(n_components=keep_dim)
    pca.fit(A_flat)
    # con = pca.get_covariance()
    # a, b = np.linalg.eig(con)
    A_flat_tran = pca.transform(A_flat)
    A_flat_restore = pca.inverse_transform(A_flat_tran)
    for f in range(fea_dim):
        A_flat_restore[:, f] += fea_means[f]
    A_restore = None
    v_restore = None
    for r in range(num_samples):
        if v_restore is None:
            v_restore = np.reshape(A_flat_restore[r], block)
        else:
            v_restore = np.hstack((v_restore, np.reshape(A_flat_restore[r], block)))
        if (r + 1) % shape[0] == 0:
            if A_restore is None:
                A_restore = v_restore
            else:
                A_restore = np.vstack((A_restore, v_restore))
            v_restore = None
    return A_restore

